import re
import process_data
import math
from pathlib import Path
from typing import Annotated, Union, Tuple, Dict
import json
import demjson
import random
import copy

# import typer
# Typer 是一个用于在 Python 中创建命令行界面（CLI）应用程序的库。
from peft import AutoPeftModelForCausalLM, PeftModelForCausalLM
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    PreTrainedModel,
    PreTrainedTokenizer,
    PreTrainedTokenizerFast,
    GenerationConfig,
)

import argparse
import os

from process_data import process_data

ModelType = Union[PreTrainedModel, PeftModelForCausalLM]
TokenizerType = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]


# app = typer.Typer(pretty_exceptions_show_locals=False)


class MyLLM:
    def __init__(self, model_dir):
        self.model_dir = model_dir
        self.model, self.tokenizer = self.load_model_and_tokenizer(model_dir)

    @staticmethod
    def load_model_and_tokenizer(
            model_dir: Union[str, Path], trust_remote_code: bool = True
    ) -> tuple[ModelType, TokenizerType]:
        model_dir = Path(model_dir).expanduser().resolve()
        if (model_dir / 'adapter_config.json').exists():
            model = AutoPeftModelForCausalLM.from_pretrained(
                model_dir, trust_remote_code=trust_remote_code, device_map='auto'
            )
            tokenizer_dir = model.peft_config['default'].base_model_name_or_path
        else:
            model = AutoModelForCausalLM.from_pretrained(
                model_dir, trust_remote_code=trust_remote_code, device_map='auto'
            )
            tokenizer_dir = model_dir
        tokenizer = AutoTokenizer.from_pretrained(
            tokenizer_dir, trust_remote_code=trust_remote_code, encode_special_tokens=True, use_fast=False
        )
        return model, tokenizer

    def respond(self, messages):
        inputs = self.tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_tensors="pt"
        ).to(self.model.device)
        generate_kwargs = {
            "input_ids": inputs,
            "max_new_tokens": 1024,
            "do_sample": True,
            "top_p": 0.8,
            "temperature": 0.8,
            "repetition_penalty": 1.2,
            "eos_token_id": self.model.config.eos_token_id,
        }
        outputs = self.model.generate(**generate_kwargs)
        response = self.tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True).strip()
        print("=========")
        print(response)
        return response

    def multi_respond(self, messages, respond_times):
        results = []
        for _ in range(respond_times):
            results.append(self.respond(messages))
        return results

    @staticmethod
    def build_message_from_template(user_prompt, assistant_prompt=None, system_prompt=None):
        if system_prompt is None:
            system_prompt = "你是一位资深中医，并且对中医古籍非常了解"
        messages = [
            {
                "role": "system", "content": system_prompt,
            },
            {
                "role": "user", "content": user_prompt,
            },
        ]
        if assistant_prompt is not None:
            messages.append({"role": "assistant", "content": assistant_prompt})
        return messages

    @staticmethod
    def correct_response_format(response: str):
        quote_pat = re.compile(r'"药物剂量":\s?\{[^{}]+\}')
        response_match = re.search(quote_pat, response + "}")
        if response_match:
            possible_json = "{" + response[response_match.start():response_match.end()] + "}"
            try:
                json_result = demjson.decode(possible_json)
                return json_result, possible_json
            except demjson.JSONDecodeError:
                return possible_json, possible_json
        json_pat = re.compile(r"json\s?\{[^{}]+\}")
        response_match = re.search(json_pat, response)
        if response_match:
            json_response = response[response_match.start():response_match.end()]
            possible_json = json_response[4:]
            try:
                json_result = demjson.decode(possible_json)
                return {"药物剂量": json_result}, possible_json
            except demjson.JSONDecodeError:
                return possible_json, possible_json

    @staticmethod
    def pre_correct_response_format(response: str):
        response = re.sub(r'//.+?\n', '\n', response)  # 去掉json中可能存在的注释
        response = re.sub(r'\s|\n|\(|\)', '', response)
        response = re.sub(r',(\"\",)+', ',', response)  # 需要注意可能有多个空的
        response = re.sub(r',\"\"\}', '}', response)
        replace_mapping_list = {"herbal_dosage": "药物剂量",
                                "组成": "药物剂量",
                                "药方成分": "药物剂量",
                                "成分": "药物剂量",
                                "．": ".",
                                "`": ""}
        for key, value in replace_mapping_list.items():
            response = response.replace(key, value)
        return response

    def parse_result(self, response, logger=None):
        response = MyLLM.pre_correct_response_format(response)
        quote_pat = re.compile(r'"药物剂量":?\{[^{}]+\}')
        response_match = re.search(quote_pat, response)
        if not response_match:
            json_result, possible_json = self.correct_response_format(response)
            for _ in range(3):
                if json_result == possible_json:
                    modified_json = self.modify_possible_json(possible_json)
                    if logger is not None:
                        logger.info("modified json %s" % modified_json)
                    modified_match = re.search(quote_pat, modified_json)
                    if modified_match:
                        # 经过llm修正格式成功，可以转成json
                        break
                else:
                    # 说明上面的correct_response_format成功，可以返回正确结果
                    return json_result, possible_json
        possible_json = "{" + response[response_match.start():response_match.end()] + "}"
        # demjson对于额外的\n不容易出错
        try:
            json_result = demjson.decode(possible_json)
        except demjson.JSONDecodeError:
            json_result = possible_json
            if logger is not None:
                logger.info("json dem error %s" % possible_json)
        return json_result, possible_json

    @staticmethod
    def flat_dict_to_str(input: dict):
        assert type(input) is dict, input
        result = ""
        for k in input:
            if input[k] is None:
                # 主治有空的
                continue
            result += k + "："
            if type(input[k]) is str:
                result += input[k] + "。"
            elif type(input[k]) is list:
                result += "，".join(input[k]) + "。"
            elif type(input[k]) is dict:
                result += "，".join([str(ik) + "：" + str(input[k][ik]) for ik in input[k]])
            else:
                print(type(input[k]), input)
                raise Exception("Unexpected type")
        return result

    def modify_possible_json(self, possible_json: str):
        user_prompt = """任务描述：请根据json格式要求，在保留正确信息的情况下，将其改造成正确的json数据。要求json的外层dict中包含"药物剂量" 输入: %s。""" % possible_json
        messages = self.build_message_from_template(user_prompt)
        response = self.respond(messages)
        return response

    def self_consistent(self, multi_responses, strategy="vote"):
        """
        对多个response进行投票
        先对每个response进行解析，看是否解析成功
        然后对每个response进行投票，取票数最多的response
        :param multi_responses: 多个response
        :param meds: 药物名
        :param strategy: 投票策略，投票或者取平均
        """
        predict_results = []
        for response in multi_responses:
            predict_result, possible_json = self.parse_result(response)
            if predict_result == possible_json:
                # 解析失败
                continue
            else:
                predict_results.append(predict_result)
        vote_result = {}
        for predict_result in predict_results:
            for med in predict_result:
                if med not in vote_result:
                    vote_result[med] = {}
                if predict_result[med] not in vote_result[med]:
                    vote_result[med][predict_result[med]] = 0
                vote_result[med][predict_result[med]] += 1
        average_flag = False
        if strategy == "average":
            average_result = {}
            try:
                for med in vote_result:
                    dose = max(vote_result[med], key=vote_result[med].get)
                    _, unit = extract_num_and_unit(dose)
                    tem = {extract_num_and_unit(dose)[0]: vote_result[med][dose] for dose in
                           vote_result[med]}
                    average_result[med] = str(sum([dose * tem[dose] for dose in tem]) / sum(
                        tem.values())) + unit
                return average_result
            except:
                print("average error")
                average_flag = True
        if strategy == "vote" or average_flag:
            final_result = {}
            for med in vote_result:
                final_result[med] = max(vote_result[med], key=vote_result[med].get)
                # 从 vote_result 字典中找出每个药物 med 得票最多的剂量，并将其作为最终结果 final_result 中该药物的剂量
            return final_result

    def respond_with_rate(self, ):
        # 目前的rate实现的是比较两个方剂中的哪一个和原方剂更接近
        prompt = """主治：%s，原方剂量：%s。一个智能体认为转换后的剂量是：%s，另一个智能体认为转换后的剂量是：%s。""" % (
            self.user_input["主治"], self.user_input["药物剂量"], self.response)


def build_predict_prompt(input_type, user_input: dict, example: dict = None):
    # user_prompt = """任务描述：请考虑方剂的名称、组成和主治，使用json格式，以克为单位补全输入文本中方剂中每一味药物的剂量。格式使用字典形式如 药物：剂量。输入: %s。""" % MyLLM.flat_dict_to_str(user_input)
    example_prompt = ""
    if example is not None:
        example_prompt = "例如: 输入：%s；输出：%s" % (
            json.dumps(example["user"], ensure_ascii=False), json.dumps(example["assistant"], ensure_ascii=False))
    if input_type == "modern":
        user_prompt = """任务描述：请考虑方剂的名称、组成和主治，使用json格式，以克为单位补全输入文本中方剂中每一味药物的剂量。格式使用字典形式如 {药物：剂量}。字典中不可以出现空键。 %s\n 输入: %s。""" % (
            example_prompt, json.dumps(user_input, ensure_ascii=False))
    elif input_type == "traditional":
        user_prompt = """任务描述：请考虑方剂的名称、组成和主治，使用json格式，使用古代常用单位补全输入文本中方剂中每一味药物的剂量，如两、钱等。格式使用字典形式如 {药物：剂量}。字典中不可以出现空键。 %s\n 输入: %s。""" % (
            example_prompt, json.dumps(user_input, ensure_ascii=False))
    else:
        raise Exception("Unexpected input type")
    return user_prompt


def build_convert_prompt(input_type, user_input: str, input_med_dose, example: dict = None):
    example_prompt = ""
    if example is not None:
        # example_prompt = "例子: 输入：%s；输出：%s" % (MyLLM.flat_dict_to_str(example["user"]), MyLLM.flat_dict_to_str(example["assistant"]))
        era = "古代" if input_type == "traditional" else "现代"
        example_prompt = "例如: 输入：主治：%s，%s剂量：%s ；输出：%s" % (
            example["主治"], era, json.dumps(example["med_dose"], ensure_ascii=False),
            json.dumps(example["assistant"], ensure_ascii=False))
    if input_type == "traditional":  # from traditional to modern
        user_prompt = """任务描述：请考虑方剂的名称、主治、组成，使用json格式，对输入文本中方剂中每一味药物将古代单位表示的剂量转换为以克为单位的剂量。格式使用字典形式如 {药物：剂量}。字典中不可以出现空键。 %s\n 主治: %s。\n古代药物剂量： %s""" % (
            example_prompt, user_input, json.dumps(input_med_dose, ensure_ascii=False))
    elif input_type == "modern":  # from modern to traditional
        user_prompt = """任务描述：请考虑方剂的名称、组成和主治，使用json格式，对输入文本中方剂中每一味药物将现代以克为单位表示的剂量转换为古代常用单位的剂量，如两、钱等。格式使用字典形式如 {药物：剂量}。字典中不可以出现空键。
                       "%s\n主治: %s。 \n现代药物剂量： %s""" % (
            example_prompt, user_input, json.dumps(input_med_dose, ensure_ascii=False))
    else:
        raise Exception("Unexpected input type")
    return user_prompt


def build_extract_prompt(user_input: dict, example: dict = None):
    user_prompt = """请将下面方剂中的药物和对应剂量提取出来，忽略其中的制法信息，并将剂量中的数量信息以阿拉伯数字表示，以json的字典格式返回，如 {药物：剂量}。需要注意三七、五倍子、五味子等包含数字的药物中的数字不做转换。如果没有剂量描述则以None代替。输入：%s""" % user_input
    return user_prompt


def build_rating_prompt(mode, cure_text, rate_dose, rate_dose_other=None, original_dose=None):
    if mode == "convert":
        assert original_dose is not None
        user_prompt = """请考虑给出的原文本中方剂中各个药物成分的剂量，以及所对应的病症。从药物出现是否完整、药物所占总体比例、药物之间的相对剂量关系等方面，判断结构化json数据中药物剂量是否合理。并给出0~1之间的得分。原文：%s，原药物剂量：%s，json结构方剂：%s""" % (
            cure_text, original_dose, rate_dose)
    elif mode == "predict":
        user_prompt = """请考虑给出的原文本中的病症和方剂的药物组成。从药物出现是否完整、药物所占总体比例、药物之间的相对剂量关系等方面，判断结构化json数据中药物剂量是否合理。并给出0~1之间的得分。原文：%s，json结构方剂：%s""" % (
            cure_text, rate_dose)
    elif mode == "compare":
        assert rate_dose_other is not None
        user_prompt = """请考虑给出的原文本中的病症和方剂的药物组成。从药物出现是否完整、药物所占总体比例、药物之间的相对剂量关系等方面，判断下列结构化json数据药物剂量哪个更加合理。并给出0~1之间的置信度。请比较两个智能体认为的转换后的剂量，哪个更接近原剂量，返回一个0到1的小数，0表示第一个智能体认为的转换后的剂量更接近原剂量，1表示第二个智能体认为的转换后的剂量更接近原剂量。原文：%s，方剂1：%s，方剂2：%s""" % (
            cure_text, rate_dose, rate_dose_other)
    elif mode == "contrastive":
        user_prompt = """请考虑给出的原文本中的病症和方剂的药物剂量。从药物出现是否完整、药物所占总体比例、药物之间的相对剂量关系等方面，判断下列结构化json数据药物剂量哪个更加合理。并给出0~1之间的置信度。请比较两个智能体认为的转换后的剂量，哪个更接近原剂量，返回一个0到1的小数，0表示第一个智能体认为的转换后的剂量更接近原剂量，1表示第二个智能体认为的转换后的剂量更接近原剂量。原文：%s，原剂量：%s，方剂1：%s，方剂2：%s""" % (
            cure_text, original_dose, rate_dose, rate_dose_other)
    else:
        raise Exception("Unexpected mode")
    return user_prompt


def modify_assistant(assistant_record: dict):
    """
    如果剂量中不含数字则置为None，不进行预测
    :param assistant_record:药物剂量的词典
    :return:
    """
    # print(tokenizer.special_tokens_map)
    num_pat = re.compile(r"[0-9\.一二三四五六七八九十百]+")
    modified = {}
    for med, dose in assistant_record:
        if not re.search(num_pat, dose):
            modified[med] = None
        else:
            modified[med] = dose
    return modified


def extract_num_and_unit(dose):
    if dose:
        dose = dose.replace("．", ".")
        num_pat = re.compile(r"[0-9\.]+")  # 这里应该不需要考虑中文的数量表示，因为前面处理数据时已经转换成阿拉伯数字了
        match = re.search(num_pat, dose)
        if match:
            return dose[match.start():match.end()], dose[match.end():]
    else:
        return None, None


def measure_dose(predict, gold):
    """
    衡量剂量的差别，目前做法会匹配第一个比如3-5只会匹配到3
    :param predict:
    :param gold:
    :return:
    """
    score = 0
    for med in gold["药物剂量"]:
        possible_gold, _ = extract_num_and_unit(gold["药物剂量"][med])
        if possible_gold:
            gold_dose = float(possible_gold)
        if med in predict["药物剂量"]:
            possible_dose, _ = extract_num_and_unit(predict["药物剂量"][med])
            if possible_dose and possible_gold:
                predict_dose = float(possible_dose)
                score += min(math.fabs(gold_dose - predict_dose), gold_dose) / gold_dose
    return score / len(gold["药物剂量"])


def measure_ratio(predict, gold):
    score = 0
    for med in gold:
        gold_dose = gold[med]
        if med in predict:
            predict_dose = predict[med]
            score += min(math.fabs(gold_dose - predict_dose), gold_dose) / gold_dose
    return score / len(gold)


def med_fidelity(predict, gold):
    pred_set = set(predict["药物剂量"])
    gold_set = set(gold["药物剂量"])
    return len(pred_set.intersection(gold_set)) / len(pred_set)


def cal_ratio(dose_pair: dict):
    extracted_meds = {}
    for med in dose_pair["药物剂量"]:
        if dose_pair["药物剂量"][med] is not None:
            possible_dose, _ = extract_num_and_unit(dose_pair["药物剂量"][med])
            if possible_dose:
                extracted_meds[med] = float(possible_dose)
    dose_sum = sum(list(extracted_meds.values()))
    result = {med: extracted_meds[med] / dose_sum for med in extracted_meds}
    return result


def extract_dose_set(meds: dict):
    dose_dict = {}
    for med in meds:
        if meds[med] not in dose_dict and meds[med] is not None:
            dose_dict[meds[med]] = set()
        if meds[med] in dose_dict:
            dose_dict[meds[med]].add(med)
    return dose_dict


def equality_measure(predict, gold):
    pred_dose_dict = extract_dose_set(predict["药物剂量"])
    gold_dose_dict = extract_dose_set(gold["药物剂量"])
    right = 0
    for dose in pred_dose_dict:
        if dose in gold_dose_dict:
            pred_meds = pred_dose_dict[dose]
            gold_meds = gold_dose_dict[dose]
            right += len(pred_meds.intersection(gold_meds))
    return right / len(gold["药物剂量"])


def my_metric(predict, gold):
    print(gold)
    fidelity = med_fidelity(predict, gold)
    dose_score = measure_dose(predict, gold)
    dose_ratio_score = measure_ratio(cal_ratio(predict), cal_ratio(gold))
    equal_score = equality_measure(predict, gold)
    result = {"剂量": dose_score, "忠实": fidelity, "等分": equal_score, "比例": dose_ratio_score}
    return result


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_dir", type=str)
    parser.add_argument("--data_dir", type=str)
    parser.add_argument("--data_file", type=str)
    parser.add_argument("--finetune_dir", type=str)
    parser.add_argument("--modify_assistant", action="store_true")

    args = parser.parse_args()
    data = process_data.build_predict_finetune_dataset(os.path.join(args.data_dir, args.data_file), data_type="modern")
    train_data, val_data, test_data = process_data.split_finetune_data(data)
    # for fname, data in zip(["train.jsonl", "val.jsonl", "test.jsonl"], [train_data, val_data, test_data]):
    process_data.write_message_jsonl(["train.jsonl", "val.jsonl", "test.jsonl"], [data, data, data], args.finetune_dir)
